#!/usr/bin/python3
import os
import tempfile


class DotGraph(object):
  """A low level interface for working with dot."""

  def __init__(self):
    """Constructs a DotGraph."""

    self.fd = tempfile.NamedTemporaryFile()
    self.fd.write(b"digraph \"\" {\n"
                  b"graph [overlap=true, splines=\"true\", start=2]\n"
                  b"node [fontname=Arial, fontsize=\"10\"];\n"
                  b"edge  [fontname=Arial, fontsize=8, arrowsize=0.5]\n")

    self._finished = False
    self._id_count = 0

  def add_pt(self, label, color="dodgerblue", mouseover=None):
    """Adds a point to the DotGraph.

    Args:
      label: string, The label to render on the point.
      color: string, The color of the point.
      mouserover: string, The mouseover text.

    Returns:
      unique ID for the point.
    """
    self._id_count += 1
    if mouseover == '""':
      mouseover = '"%d"' % self._id_count

    if not mouseover:
      mouseover = self._id_count
    self.fd.write(bytes(str(mouseover) +
       "[label=\"" + label + "\" color=\"" + color + "\" ];\n", "utf8"))
    return mouseover

  def line(self, start, end, label=None, both_directions=False):
    """Draws a line from one node to the next node.
    
    Args:
      start: string, The node unique id returned by add_pt to draw a line from.
      end: string, The node unique id returned by add_pt to draw a line to.
      both_directions: bool, If true, the arrow will go both ways.
    """
    args = []
    if label:
      args.append('[label="' + label + '"]')
    if both_directions:
      args.append('[dir="both"]')
    self.fd.write(bytes(str(start) + ' -> ' + str(end) + ''.join(args) + ';\n', 'utf8'))

  def _finish(self):
    """Finishes the graph and flushes it to the file."""
    if not self._finished:
      self.fd.write(b"}\n")
      self.fd.flush()
      self._finished = True

  def show(self):
    """Renders the graph, and loads it in Chrome."""
    self._finish()
    self.out = tempfile.NamedTemporaryFile()
    os.system("dot " + self.fd.name + " -Tsvg > " + self.out.name)
    os.system("google-chrome " + self.out.name)

  def draw_to(self, fname):
    """Renders the graph to the specified file.

    Args:
      fname: string, The filename to render to.
    """
    self._finish()
    os.system("dot " + self.fd.name + " -Tsvg > " + fname)


class DrawGraph(object):
  """Object that draws rules, the associated filesets, and the connections."""

  def __init__(self):
    # The underlying dot interface.
    self._out_graph = DotGraph()
    # A unique ID for the nodes.
    self._fileset_id_count = 0
    # If a fileset or rule is in this cache, it will be mapped to the
    # associated node ID used by DotGraph.
    self._cache = dict()

  def _add_fileset(self, fileset):
    """Adds a fileset and all rule and fileset dependencies to the graph
    
    Args:
      fileset: FileSet, the fileset.
    """
    if fileset in self._cache:
      return self._cache[fileset]

    # Generate a node mouseover.
    self._fileset_id_count += 1
    mouseover = '\"' + str(self._fileset_id_count) + ' ' + str(fileset) + '\"'

    # Add the point
    files = fileset.files
    if len(files) == 0:
      files = ['----']
    node = self._out_graph.add_pt(', '.join(files), mouseover=mouseover)
    self._cache[fileset] = node

    # Add all the dependencies to the graph.
    for dep in fileset.dependencies:
      depnode = self._add_fileset(dep)
      if depnode:
        self._out_graph.line(depnode, node)

    # Make sure the parent rule has also been added.
    rule = fileset.gen_rule
    if rule:
      self.add_rule(fileset.gen_rule)
    else:
      #  No parrent rule is bad, draw the missing rule as a red circle.
      key = fileset.key
      rulenode = self._out_graph.add_pt(
          key.genrule_key.full_name(), color='red',
          mouseover='\"' + str(key.genrule_key) + '\"')
      self._out_graph.line(rulenode, node, label=key.bucket_name + '.' + key.name)
    return node

  def add_rule(self, rule):
    """Adds a rule to the graph, including all dependent rules and filesets.

    Args:
      rule: cache_view.GenRule, The rule to graph.
    """
    # Only draw a rule once.
    if rule in self._cache:
      return self._cache[rule]

    # Change the color if the rule is valid.
    color = 'green' if rule.get_valid() else 'darkorange'
    node = self._out_graph.add_pt(
        str(rule), color=color,
        mouseover='\"' + ' '.join(rule.get_cmd()) + '\"')

    # Declare the rule drawn, and move on.
    self._cache[rule] = node

    self._add_all_filesets(node, rule.inp, 'inp', inp=True)
    self._add_all_filesets(node, rule.out, 'out', output=True)
    self._add_all_filesets(node, rule.meta, 'meta', inp=True, output=True)

  def _add_all_filesets(self, node, file_bucket, name, inp=False, output=False):
    """Adds all the filesets in a file bucket to the dot graph.

    Args:
      file_bucket: FileBucket, A bucket of all the filesets to add to the graph.
      name: string, The name of the file_bucket.
      node: The unique node id.
      output: bool, If true, the fileset is an output and the arrows should be
        drawn accordingly
      inp: bool, If true, the fileset is an input and the arrows should be
        drawn accordingly.  output and inp can both be specified.
    """
    for tag in file_bucket.getKeys():
      fileset = getattr(file_bucket, tag)
      fnode = self._add_fileset(fileset)
      if fnode:
        if output and not inp:
          start, end = node, fnode
          self._out_graph.line(start, end, label=name + tag)
        elif inp and not output:
          start, end = fnode, node
          self._out_graph.line(start, end, label=name + tag)
        elif inp and output:
          self._out_graph.line(fnode, node, label=name + tag,
                               both_directions=True)

  def show(self):
    """Renders the graph and opens it in Chrome."""
    self._out_graph.show()

  def draw_to(self, fname):
    """Renders the graph to the specified file.

    Args:
      fname: string, The filename to render to.
    """
    self._out_graph.draw_to(fname)
